{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 深度概率编程之贝叶斯层转换\n",
    "\n",
    "## 概述\n",
    "\n",
    "对于不熟悉贝叶斯模型的深度神经网络（Deep Neural Networks，简称DNN）研究人员，MDP提供了高级APITransformToBNN，支持深度神经网络（Deep Neural Networks）模型一键转换成贝叶斯神经网络（Bayes Neural Networks，后续简称BNN）模型。目前在LeNet，ResNet，MobileNet，VGG等模型上验证了API的通用性。本例将会介绍如何使用transforms模块中的TransformToBNNAPI实现DNN一键转换成BNN。  \n",
    "\n",
    "整体流程如下：\n",
    "\n",
    "1. 定义DNN模型；\n",
    "2. 定义损失函数和优化器；\n",
    "3. 转换整个模型；\n",
    "\n",
    "> 本例适用于GPU和Ascend环境。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义DNN模型\n",
    "\n",
    "本例用到的深度神经网络(DNN)模型为LeNet5，定义完成后，打印其神经层的名称。由于转换层面上主要卷积层和池化层，本例也针对性的展示这两种计算层的转换信息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore.nn as nn\n",
    "from mindspore.common.initializer import Normal\n",
    "\n",
    "class LeNet5(nn.Cell):\n",
    "    \"\"\"Lenet network structure.\"\"\"\n",
    "    # define the operator required\n",
    "    def __init__(self, num_class=10, num_channel=1):\n",
    "        super(LeNet5, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')\n",
    "        self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))\n",
    "        self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))\n",
    "        self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))\n",
    "        self.relu = nn.ReLU()\n",
    "        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.flatten = nn.Flatten()\n",
    "\n",
    "    # use the preceding operators to construct networks\n",
    "    def construct(self, x):\n",
    "        x = self.max_pool2d(self.relu(self.conv1(x)))\n",
    "        x = self.max_pool2d(self.relu(self.conv2(x)))\n",
    "        x = self.flatten(x)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.fc3(x) \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于经典的DNN模型LeNet5网络卷积层有两层conv1，conv2，全连接层为3层：fc1，fc2，fc3。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义损失函数和优化器\n",
    "\n",
    "本例中使用损失函数为交叉熵损失函数`nn.SoftmaxCrossEntropyWithLogits`，优化器为`Adam`函数即`nn.AdamWeightDecay`。\n",
    "\n",
    "由于需要将进行整个模型的BNN转换，所以需要将DNN网络，损失函数和优化器关联成一个完整的计算网络，即`train_network`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['conv1.weight',\n",
      " 'conv2.weight',\n",
      " 'fc1.weight',\n",
      " 'fc1.bias',\n",
      " 'fc2.weight',\n",
      " 'fc2.bias',\n",
      " 'fc3.weight',\n",
      " 'fc3.bias']\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "import numpy as np\n",
    "from mindspore.nn import WithLossCell, TrainOneStepCell\n",
    "from mindspore.nn.probability import transforms\n",
    "from mindspore import context\n",
    "\n",
    "context.set_context(mode=context.GRAPH_MODE,device_target=\"GPU\")\n",
    "\n",
    "network = LeNet5()\n",
    "lr = 0.01\n",
    "momentum = 0.9\n",
    "criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction=\"mean\")\n",
    "optimizer = nn.AdamWeightDecay(params=network.trainable_params(), learning_rate=0.0001)\n",
    "#optimizer = nn.Momentum(network.trainable_params(), lr, momentum)\n",
    "net_with_loss = WithLossCell(network, criterion)\n",
    "train_network = TrainOneStepCell(net_with_loss, optimizer)\n",
    "\n",
    "DNN_layer_name = [i.name for i in network.trainable_params()]\n",
    "pprint.pprint(DNN_layer_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上述打印信息即为当前未转换的卷积层和全连接层名称。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 转换整个模型\n",
    "\n",
    "转换整个模型使用到了`transforms`中的`TransformToBNN`API，一键转换完成后打印出模型中BNN的名称。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['conv1.weight_posterior.mean',\n",
      " 'conv1.weight_posterior.untransformed_std',\n",
      " 'conv2.weight_posterior.mean',\n",
      " 'conv2.weight_posterior.untransformed_std',\n",
      " 'fc1.weight_posterior.mean',\n",
      " 'fc1.weight_posterior.untransformed_std',\n",
      " 'fc1.bias_posterior.mean',\n",
      " 'fc1.bias_posterior.untransformed_std',\n",
      " 'fc2.weight_posterior.mean',\n",
      " 'fc2.weight_posterior.untransformed_std',\n",
      " 'fc2.bias_posterior.mean',\n",
      " 'fc2.bias_posterior.untransformed_std',\n",
      " 'fc3.weight_posterior.mean',\n",
      " 'fc3.weight_posterior.untransformed_std',\n",
      " 'fc3.bias_posterior.mean',\n",
      " 'fc3.bias_posterior.untransformed_std']\n"
     ]
    }
   ],
   "source": [
    "bnn_transformer = transforms.TransformToBNN(train_network, 60000, 0.000001)\n",
    "train_bnn_network = bnn_transformer.transform_to_bnn_model()\n",
    "BNN_layer_name =[i.name for i in network.trainable_params()]\n",
    "\n",
    "pprint.pprint(BNN_layer_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上述打印信息即整体转换成贝叶斯网络（BNN）后的卷积层和全连名称。  \n",
    "\n",
    "如果想要进行单个贝叶斯层的转换可参考[官网教程](https://www.mindspore.cn/tutorial/training/zh-CN/master/advanced_use/apply_deep_probability_programming.html#id16)。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore-1.0.1",
   "language": "python",
   "name": "mindspore-1.0.1"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
